#include <stdlib.h>
#include "bstree.h"
#include <stdio.h>

void left_rotate(ptnode t, ptnode x)
{
  ptnode y = x->right;

  x->right = y->left;
  y->left->p = x;
  y->p = x->p;
  if (x->p == nil)
    newtreeroot(t, y);
  else if (x == x->p->left)
    x->p->left = y;
  else
    x->p->right = y;
  y->left = x;
  x->p = y;
  y->size = x->size;
  x->size = x->left->size + x->right->size + 1;
}

void right_rotate(ptnode t, ptnode y)
{
  ptnode x = y->left;

  y->left = x->right;
  x->right->p = y;
  x->p = y->p;
  if (y->p == nil)
    newtreeroot(t, x);
  else if (y == y->p->left)
    y->p->left = x;
  else
    y->p->right = x;
  x->right = y;
  y->p = x;
  x->size = y->size;
  y->size = y->left->size + y->right->size + 1;
}

void rb_insert_fixup(ptnode t, ptnode z)
{
  ptnode y;

  while (z->p->color == 'R') {
    if (z->p == z->p->p->left) {
      y = z->p->p->right;
      if (y->color == 'R') {
	z->p->color = y->color = 'B';
	z->p->p->color = 'R';
	z = z->p->p;
      }
      else {
	if (z == z->p->right) {
	  z = z->p;
	  left_rotate(t, z);
      	}
	z->p->color = 'B';
	z->p->p->color = 'R';
	right_rotate(t, z->p->p);
      }
    }
    else {
      y = z->p->p->left;
      if (y->color == 'R') {
	z->p->color = y->color = 'B';
	z->p->p->color = 'R';
	z = z->p->p;
      }
      else {
	if (z == z->p->left) {
	  z = z->p;
	  right_rotate(t, z);
	}
	z->p->color = 'B';
	z->p->p->color = 'R';
	left_rotate(t, z->p->p);
      }
    }
  }
  treeroot(t)->color = 'B';
}

void rb_insert(ptnode t, ptnode z)
{
  ptnode x, y;
  
  x = treeroot(t);
  y = nil;
  while (x != nil) {
    y = x;
    if (z->key < x->key)
      x = x->left;
    else 
      x = x->right;
  }
  z->p = y;
  if (y == nil)
    newtreeroot(t, z);
  else if (z->key < y->key)
    y->left = z;
  else
    y->right = z;
  z->color = 'R';
  z->size = 1;		      /* maintain size from z up to the root*/
  x = z->p;			
  while (x != nil) {
    (x->size)++;
    x = x->p;
  }
  rb_insert_fixup(t, z);
} 

void rb_delete_fixup(ptnode t, ptnode x)
{
  ptnode w;

  while (x != treeroot(t) && x->color == 'B') {
    if (x == x->p->left) {
      w = x->p->right;
      if (w->color == 'R') {
	w->color = 'B';
	x->p->color = 'R';
	left_rotate(t, x->p);
	w = x->p->right;
      }
      if (w->left->color == w->right->color == 'B') {
	w->color = 'R';
	x = x->p;
      }
      else {
	if (w->right->color == 'B') {
	  w->left->color = 'B';
	  w->color = 'R';
	  right_rotate(t, w);
	  w = x->p->right;
	}
	w->color = x->p->color;
	x->p->color = 'B';
	w->right->color = 'B';
	left_rotate(t, x->p);
	x = t;
      }
    }
    else {
      w = x->p->left;
      if (w->color == 'R') {
	w->color = 'B';
	x->p->color = 'R';
	right_rotate(t, x->p);
	w = x->p->left;
      }
      if (w->right->color == w->left->color == 'B') {
	w->color = 'R';
	x = x->p;
      }
      else {
	if (w->left->color == 'B') {
	  w->right->color = 'B';
	  w->color = 'R';
	  left_rotate(t, w);
	  w = x->p->left;
	}
	w->color = x->p->color;
	x->p->color = 'B';
	w->left->color = 'B';
	right_rotate(t, x->p);
	x = treeroot(t);
      }
    }
  }
  x->color = 'B';
}

ptnode rb_delete(ptnode t, ptnode z)
{
  ptnode x, y, w;

  if (z->left == nil || z->right == nil)
    y = z;
  else
    y = tree_successor(z);
  if (y->left != nil)
    x = y->left;
  else
    x = y->right;
  x->p = y->p;
  if (y->p == nil)		/* y is root */
    newtreeroot(t, x);
  else if (y == y->p->left)
    y->p->left = x;
  else
    y->p->right = x;
  if (y != z)
    z->key = y->key;
  w = y->p;			
  while (w != nil) {
    (w->size)--;
    w = w->p;
  }
  if (y->color == 'B')
    rb_delete_fixup(t, x);
  return y;
} 

long blackheight(ptnode x)
{
  long h = 0;

  while (x->left != nil) {	/* x->right will also be ok */
    x = x->left;
    if (x->color == 'B')
      h++;
  }
  h++;

  return h;
}

long os_key_rank(ptnode x, long k) /* x: root[T] */
{
  if (x->key == k)
    return x->left->size + 1;
  else if (k < x->key)
    return os_key_rank(x->left, k);
  else
    return x->left->size + 1 + os_key_rank(x->right, k);
}



